#!/usr/bin/env python3
# -*- coding:utf-8 -*-

"""
【节点名称】：
    RTSPStreamerNode
【依赖项安装】：

    1 安装 GStreamer-1.0及相关插件:
        sudo apt-get install libgstreamer1.0-0 gstreamer1.0-plugins-base gstreamer1.0-plugins-good gstreamer1.0-plugins-bad gstreamer1.0-plugins-ugly gstreamer1.0-libav gstreamer1.0-doc gstreamer1.0-tools gstreamer1.0-x gstreamer1.0-alsa gstreamer1.0-gl gstreamer1.0-gtk3 gstreamer1.0-qt5 gstreamer1.0-pulseaudio
    2 安装 RTSP 服务器:
        sudo apt-get install libglib2.0-dev libgstrtspserver-1.0-dev gstreamer1.0-rtsp
    3 安装 OPENCV-PYTHON
        pip install opencv-python
    3 安装 spirems及相关依赖
        pip install spirems
        git clone -b spirecv2-dds https://gitee.com/amovlab/SpireCV.git
        cd SpireCV && pip install -e .
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
    sensor_msgs::CameraCalibration (输入相机参数)
【发布类型】：
    无
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号,如果需要长期后台执行,建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
【节点参数】：

【备注】：
    无
"""



import threading
import time
import cv2
import os
import json
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg
import argparse
from spirecv.base.BaseNode import BaseNode

# import necessary argumnets 
import gi

# import required library like Gstreamer and GstreamerRtspServer
gi.require_version('Gst', '1.0')
gi.require_version('GstRtspServer', '1.0')
from gi.repository import Gst, GstRtspServer, GObject, GLib


# Sensor Factory class which inherits the GstRtspServer base class and add
# properties to it.
class SensorFactory(GstRtspServer.RTSPMediaFactory, threading.Thread, BaseNode):
    def __init__(self,
                 job_name: str,
                 ip: str = '127.0.0.1',
                 port: int = 9094,
                 param_dict_or_file = None):
        
        super(SensorFactory, self).__init__()
        
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            "RTSPStreamerNode",
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file
        )
        
        self.image_sms_url = self.get_param("image_sms_url", "/live/sensor/image_raw")
        self.resize = self.get_param("resize", [640, 480])
        self.fps = self.get_param("fps", 30)
        self.port = self.get_param("port", 8554)
        self.stream_uri = self.get_param("stream_uri", "/live")

        self._image_reader = Subscriber(
            self.image_sms_url, 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port
        )

        self.image_queue = Queue()
        self.queue_pool.append(self.image_queue)

        self.start()
        
        self.number_frames = 0
        self.duration = 1 / self.fps * Gst.SECOND  # duration of a frame in nanoseconds
        self.launch_string = 'appsrc name=source is-live=true block=true format=GST_FORMAT_TIME ' \
                             'caps=video/x-raw,format=BGR,width={},height={},framerate={}/1 ' \
                             '! videoconvert ! nvvidconv ! video/x-raw(memory:NVMM) ' \
                             '! nvv4l2h264enc  insert-sps-pps=true bitrate= 10240000 ! h264parse' \
                             '! rtph264pay config-interval=1 name=pay0 pt=96' \
                             .format(self.resize[0], self.resize[1], self.fps)
                
    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()

    def image_callback(self, msg):
        self.image_queue.put(msg)
        
    # 从SpireMS捕获视频源并将其推送到流媒体缓冲区。
    def on_need_data(self, src, length):
        if self.is_running:
            msg1 = self.image_queue.get(block=True)
            
            frame = sms2cvimg(msg1)
            
            # 最好改变相机的分辨率
            # 而不是改变图像形状，因为它会影响图像质量。
            frame = cv2.resize(frame, self.resize, \
                        interpolation = cv2.INTER_LINEAR)
            data = frame.tostring()
            buf = Gst.Buffer.new_allocate(None, len(data), None)
            buf.fill(0, data)
            buf.duration = self.duration
            timestamp = self.number_frames * self.duration
            buf.pts = buf.dts = int(timestamp)
            buf.offset = timestamp
            self.number_frames += 1
            retval = src.emit('push-buffer', buf)
            print('pushed buffer, frame {}, duration {} ns, durations {} s'.format(self.number_frames,
                                                                                        self.duration,
                                                                                        self.duration / Gst.SECOND))
            if retval != Gst.FlowReturn.OK:
                print(retval)
    # attach the launch string to the override method
    def do_create_element(self, url):
        return Gst.parse_launch(self.launch_string)
    
    # attaching the source element to the rtsp media
    def do_configure(self, rtsp_media):
        self.number_frames = 0
        appsrc = rtsp_media.get_element().get_child_by_name('source')
        appsrc.connect('need-data', self.on_need_data)

# Rtsp server implementation where we attach the factory sensor with the stream uri
class GstServer(GstRtspServer.RTSPServer):
    def __init__(self, job_name, param_dict_or_file):
        super(GstServer, self).__init__()
        self.factory = SensorFactory(job_name=job_name, param_dict_or_file = param_dict_or_file)
        self.factory.set_shared(True)
        self.set_service(str(self.factory.port))
        self.get_mount_points().add_factory(self.factory.stream_uri, self.factory)
        self.attach(None)

def RTSPStreamerNode(job_name, param_dict_or_file):
    Gst.init(None)
    server = GstServer(job_name, param_dict_or_file)
    loop = GLib.MainLoop()
    loop.run()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'input_topic',
        help='SpireCV2 Image Topic')
    parser.add_argument(
        'config',
        help='SpireCV2 config')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args = parser.parse_args()
    print("input-topic:", args.input_topic)
    print("job-name:", args.job_name)
    print("config:", args.config)
    #dataloader = CameraNode("live", param_dict_or_file=parameter_file)
    RTSPStreamerNode(job_name=args.job_name,  param_dict_or_file=args.config)
    

